#ifndef _PE_HEADER_H_
#define _PE_HEADER_H_

#include "bitness.h"

#include <cstdint>
#include <string>
#include <type_traits>

namespace analyzer::windows
{

const auto image_number_of_directory_entries = 16;

const auto image_directory_entry_export = 0;
const auto image_directory_entry_import = 1;
const auto image_directory_entry_resource = 2;
const auto image_directory_entry_exception = 3;
const auto image_directory_entry_security = 4;
const auto image_directory_entry_baserelog = 5;
const auto image_directory_entry_debug = 6;
const auto image_directory_entry_architecture = 7;
const auto image_directory_entry_globalptr = 8;
const auto image_directory_entry_tls = 9;
const auto image_directory_entry_load_config = 10;
const auto image_directory_entry_bound_import = 11;
const auto image_directory_entry_iat = 12;
const auto image_directory_entry_delay_import = 13;
const auto image_directory_entry_com_descriptor = 14;

const auto image_nt_optional_hdr32_magic = 0x10B;
const auto image_nt_optional_hdr64_magic = 0x20B;
const auto image_rom_optional_hdr_magic = 0x107;

const auto image_sizeof_short_name = 8;

struct alignas(1) image_file_header
{
    uint16_t machine;
    uint16_t number_of_sections;
    uint32_t time_date_stamp;
    uint32_t pointer_to_symbol_table;
    uint32_t number_of_symbols;
    uint16_t size_of_optional_header;
    uint16_t characteristics;
};

struct alignas(1) image_data_directory
{
    uint32_t virtual_address;
    uint32_t size;
};

struct alignas(1) image_optional_header32
{
    uint16_t magic;
    uint8_t major_linker_version;
    uint8_t minor_linker_version;
    uint32_t size_of_code;
    uint32_t size_of_initialized_data;
    uint32_t size_of_uninitialized_data;
    uint32_t address_of_entry_point;
    uint32_t base_of_code;
    uint32_t base_of_data;
    uint32_t image_base;
    uint32_t section_alignment;
    uint32_t file_alignment;
    uint16_t major_operating_system_version;
    uint16_t minor_operating_system_version;
    uint16_t major_image_version;
    uint16_t minor_image_version;
    uint16_t major_subsystem_version;
    uint16_t minor_subsystem_version;
    uint32_t win32_version_value;
    uint32_t size_of_image;
    uint32_t size_of_headers;
    uint32_t check_sum;
    uint16_t subsystem;
    uint16_t dll_characteristics;
    uint32_t size_of_stack_reserve;
    uint32_t size_of_stack_commit;
    uint32_t size_of_heap_reserve;
    uint32_t size_of_heap_commit;
    uint32_t loader_flags;
    uint32_t number_of_rva_and_sizes;
    image_data_directory data_directory[image_number_of_directory_entries];
};

struct alignas(1) image_optional_header64
{
    uint16_t magic;
    uint8_t major_linker_version;
    uint8_t minor_linker_version;
    uint32_t size_of_code;
    uint32_t size_of_initialized_data;
    uint32_t size_of_uninitialized_data;
    uint32_t address_of_entry_point;
    uint32_t base_of_code;
    uint64_t image_base;
    uint32_t section_alignment;
    uint32_t file_alignment;
    uint16_t major_operating_system_version;
    uint16_t minor_operating_system_version;
    uint16_t major_image_version;
    uint16_t minor_image_version;
    uint16_t major_subsystem_version;
    uint16_t minor_subsystem_version;
    uint32_t win32_version_value;
    uint32_t size_of_image;
    uint32_t size_of_headers;
    uint32_t check_sum;
    uint16_t subsystem;
    uint16_t dll_characteristics;
    uint64_t size_of_stack_reserve;
    uint64_t size_of_stack_commit;
    uint64_t size_of_heap_reserve;
    uint64_t size_of_heap_commit;
    uint32_t loader_flags;
    uint32_t number_of_rva_and_sizes;
    image_data_directory data_directory[image_number_of_directory_entries];
};

template <uint8_t bitness, typename image_optional_header = std::conditional_t<bitness == x86, image_optional_header32, std::conditional_t<bitness == x64, image_optional_header64, std::false_type>>>
struct alignas(1) image_nt_header
{
    uint32_t signature;
    image_file_header file_header;
    image_optional_header optional_header;
};

struct alignas(1) image_section_header
{
    uint8_t name[image_sizeof_short_name];
    union
    {
        uint32_t physical_address;
        uint32_t virtual_size;
    } misc;
    uint32_t virtual_address;
    uint32_t size_of_raw_data;
    uint32_t pointer_to_raw_data;
    uint32_t pointer_to_relocations;
    uint32_t pointer_to_linenumbers;
    uint16_t number_of_relocations;
    uint16_t number_of_linenumbers;
    uint32_t characteristics;
};

struct alignas(1) image_import_descriptor
{
    union
    {
        uint32_t characteristics;
        uint32_t original_first_thunk;
    } dummy;
    uint32_t time_date_stamp;
    uint32_t forwarder_chain;
    uint32_t name;
    uint32_t first_thunk;
};

template <uint8_t bitness>
struct alignas(1) image_thunk_data;

template <>
struct alignas(1) image_thunk_data<x86>
{
    union
    {
        uint32_t forwarder_string;
        uint32_t function;
        uint32_t ordinal;
        uint32_t address_of_data;
    } u1;
};

template <>
struct alignas(1) image_thunk_data<x64>
{
    union
    {
        uint64_t forwarder_string;
        uint64_t function;
        uint64_t ordinal;
        uint64_t address_of_data;
    } u1;
};

struct alignas(1) image_import_by_name
{
    uint16_t hint;
    std::string name;
};

} // namespace analyzer::windows

#endif
